package com.strp.mptcp.controller;

import com.strp.mptcp.DAO.entity.FileUploadRecord;
import com.strp.mptcp.DAO.mapper.FileUploadRecordRepository;
import com.strp.mptcp.common.MPTCPMetrics;
import com.strp.mptcp.config.MPTCPConfig;
import com.strp.mptcp.socket.MultipathSocket;
import lombok.RequiredArgsConstructor;
import org.springframework.core.io.FileSystemResource;
import org.springframework.core.io.Resource;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import java.io.File;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

@RestController
@RequestMapping("/download")
@RequiredArgsConstructor
public class DownloadController {
    private final FileUploadRecordRepository recordRepository;
    private final MPTCPMetrics metrics;
    private final MPTCPConfig mptcpConfig;

    // 存储正在进行的下载
    private final Map<String, MultipathSocket> activeDownloads = new ConcurrentHashMap<>();

    /**
     * 开始文件下载
     */
    @GetMapping("/{fileHash}")
    public ResponseEntity<?> startDownload(@PathVariable String fileHash) {
        try {
            // 1. 查找文件记录
            Optional<FileUploadRecord> record = recordRepository.findByFileHash(fileHash);
            if (record.isEmpty()) {
                return ResponseEntity.notFound().build();
            }

            // 2. 创建下载目录
            File downloadDir = new File(mptcpConfig.getStorage().getDownloadDir());
            if (!downloadDir.exists()) {
                downloadDir.mkdirs();
            }

            // 3. 检查文件是否已下载完成
            File downloadedFile = new File(downloadDir, record.get().getFileName());
            if (downloadedFile.exists() && downloadedFile.length() == record.get().getFileSize()) {
                // 文件已存在且大小正确，直接返回
                return createFileResponse(downloadedFile, record.get().getFileName());
            }

            // 4. 初始化 MPTCP 连接
            List<SocketAddress> serverAddresses = mptcpConfig.getServer().getAddresses().stream()
                    .map(addr -> new InetSocketAddress(addr.getHost(), addr.getPort()))
                    .collect(Collectors.toList());

            MultipathSocket mptcpSocket = new MultipathSocket(
                    serverAddresses,
                    metrics,
                    mptcpConfig.getChunk().getSize(),
                    mptcpConfig.getChunk().getMaxRetries(),
                    mptcpConfig.getTimeout().getConnection(),
                    mptcpConfig.getTimeout().getRead(),
                    mptcpConfig.getTimeout().getWrite()
            );

            activeDownloads.put(fileHash, mptcpSocket);

            // 5. 开始下载
            mptcpSocket.downloadFile(downloadedFile, record.get().getOffset());

            // 6. 返回下载进度接口
            return ResponseEntity.ok(new DownloadResponse(fileHash, record.get().getFileName()));
        } catch (Exception e) {
            return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR)
                    .body(new DownloadResponse("Download failed: " + e.getMessage()));
        }
    }

    /**
     * 获取下载进度
     */
    @GetMapping("/progress/{fileHash}")
    public ResponseEntity<ProgressResponse> getProgress(@PathVariable String fileHash) {
        MultipathSocket socket = activeDownloads.get(fileHash);
        if (socket == null) {
            return ResponseEntity.notFound().build();
        }

        double progress = socket.getTransferProgress(fileHash);
        Map<Integer, MPTCPMetrics.SubflowStats> subflowStats = new ConcurrentHashMap<>();

        // 获取每个子流的统计信息
        List<SocketAddress> serverAddresses = mptcpConfig.getServer().getAddresses().stream()
                .map(addr -> new InetSocketAddress(addr.getHost(), addr.getPort()))
                .collect(Collectors.toList());

        for (int i = 0; i < serverAddresses.size(); i++) {
            subflowStats.put(i, socket.getSubflowStats(i));
        }

        return ResponseEntity.ok(new ProgressResponse(progress, subflowStats));
    }

    /**
     * 取消下载
     */
    @PostMapping("/cancel/{fileHash}")
    public ResponseEntity<Void> cancelDownload(@PathVariable String fileHash) {
        MultipathSocket socket = activeDownloads.remove(fileHash);
        if (socket != null) {
            socket.close();
            return ResponseEntity.ok().build();
        }
        return ResponseEntity.notFound().build();
    }

    /**
     * 恢复下载
     */
    @PostMapping("/resume/{fileHash}")
    public ResponseEntity<DownloadResponse> resumeDownload(@PathVariable String fileHash) {
        try {
            // 1. 获取文件记录
            Optional<FileUploadRecord> record = recordRepository.findByFileHash(fileHash);
            if (record.isEmpty()) {
                return ResponseEntity.notFound().build();
            }

            // 2. 检查已下载的文件
            File downloadedFile = new File(mptcpConfig.getStorage().getDownloadDir(), record.get().getFileName());
            if (!downloadedFile.exists()) {
                return ResponseEntity.badRequest()
                        .body(new DownloadResponse("Downloaded file not found"));
            }

            // 3. 创建新的 MPTCP 连接
            List<SocketAddress> serverAddresses = mptcpConfig.getServer().getAddresses().stream()
                    .map(addr -> new InetSocketAddress(addr.getHost(), addr.getPort()))
                    .collect(Collectors.toList());

            MultipathSocket mptcpSocket = new MultipathSocket(
                    serverAddresses,
                    metrics,
                    mptcpConfig.getChunk().getSize(),
                    mptcpConfig.getChunk().getMaxRetries(),
                    mptcpConfig.getTimeout().getConnection(),
                    mptcpConfig.getTimeout().getRead(),
                    mptcpConfig.getTimeout().getWrite()
            );

            activeDownloads.put(fileHash, mptcpSocket);

            // 4. 从断点继续下载
            mptcpSocket.downloadFile(downloadedFile, downloadedFile.length());

            return ResponseEntity.ok(new DownloadResponse(fileHash, record.get().getFileName()));
        } catch (Exception e) {
            return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR)
                    .body(new DownloadResponse("Resume failed: " + e.getMessage()));
        }
    }

    /**
     * 创建文件下载响应
     */
    private ResponseEntity<Resource> createFileResponse(File file, String fileName) throws IOException {
        Resource resource = new FileSystemResource(file);

        return ResponseEntity.ok()
                .contentType(MediaType.APPLICATION_OCTET_STREAM)
                .header(HttpHeaders.CONTENT_DISPOSITION, "attachment; filename=\"" + fileName + "\"")
                .header(HttpHeaders.CONTENT_LENGTH, String.valueOf(file.length()))
                .body(resource);
    }

    // DTOs
    @lombok.Data
    static class DownloadResponse {
        private final String fileHash;
        private final String fileName;
        private final String error;

        public DownloadResponse(String fileHash, String fileName) {
            this.fileHash = fileHash;
            this.fileName = fileName;
            this.error = null;
        }

        public DownloadResponse(String error) {
            this.fileHash = null;
            this.fileName = null;
            this.error = error;
        }
    }

    @lombok.Data
    static class ProgressResponse {
        private final double progress;
        private final Map<Integer, MPTCPMetrics.SubflowStats> subflowStats;
    }
} 